# coding=utf-8
##
## Author: jmdvirus@aliyun.com
##
## Create: 2019年02月12日 星期二 18时07分43秒
##

import numpy as np
import os
import gzip
import struct
import logging
import mxnet as mx
import matplotlib.pyplot as plt

logging.getLogger().setLevel(logging.DEBUG)

def read_data(label_url, image_url):
    with gzip.open(label_url) as flbl:
        magic,num = struct.unpack(">II", flbl.read(8))
        label = np.fromstring(flbl.read(), dtype = np.int8)
    with gzip.open(image_url, 'rb') as fimg:
        magic,num,rows,cols = struct.unpack(">IIII", fimg.read(16))
        image = np.fromstring(fimg.read(), dtype=np.int8)
        image = image.reshape(len(label), 1, rows, cols)
        image = image.astype(np.float32)/255.0
    return (label, image)

data_path = "/opt/data/data/av.data.mnist/"
(train_lbl, train_img) = read_data(data_path + 'train-labels-idx1-ubyte.gz', data_path + 'train-images-idx3-ubyte.gz')

(val_lbl, val_img) = read_data(data_path + 't10k-labels-idx1-ubyte.gz', data_path + 't10k-images-idx3-ubyte.gz')

batch_size = 32

train_iter = mx.io.NDArrayIter(train_img, train_lbl, batch_size, shuffle = True)
val_iter = mx.io.NDArrayIter(val_img, val_lbl, batch_size)

#for i in range(10):
#    plt.subplot(1, 10, i+1)
#    plt.imshow(train_img[i].reshape(28,28), cmap='Greys_r')
#    plt.axis('off')
#plt.show()
print('label: %s' % (train_lbl[0:10],))

data = mx.symbol.Variable('data')

flatten = mx.sym.Flatten(data = data, name = "flatten")

fc1 = mx.sym.FullyConnected(data=flatten, num_hidden=128, name="fc1")
act1 = mx.sym.Activation(data = fc1, act_type="relu", name='act1')

fc2 = mx.sym.FullyConnected(data=act1, num_hidden=64, name='fc2')
act2 = mx.sym.Activation(data = fc2, act_type = "relu", name="act2")

fc3 = mx.sym.FullyConnected(data = act2, num_hidden=10, name='fc3')

net = mx.sym.SoftmaxOutput(data=fc3, name='softmax')

shape = { "data" : (batch_size, 1, 28, 28) }
mx.viz.print_summary(symbol=net, shape=shape)

#mx.viz.plot_network(symbol=net, shape=shape).view()

module = mx.mod.Module(symbol=net, context=mx.cpu(0))

def train_mod():
    module.fit(
            train_iter,
            eval_data = val_iter,
            optimizer = 'sgd',
            optimizer_params = {'learning_rate':0.2,
                'lr_scheduler' : mx.lr_scheduler.FactorScheduler(step=60000/batch_size, factor=0.9) },
            num_epoch = 20,
            batch_end_callback = mx.callback.Speedometer(batch_size, 60000/batch_size)
            )

def save_network(mod):
    mod._symbol.save("%s-symbol.json" % "jkmnist")

def save_param(mod):
    mod.save_params("jkmnist")

def load_network(filename):
    return mx.sym.load('%s-symbol.json' % filename)

def load_param(filename):
    return mx.mod.init_params(initializer=mx.init.Load(filename))

if __name__ == "__main__":
    train_mod()
    # save_param(module)
    save_network(module)

